import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
from math import isclose
import random
from collections import Counter
import numpy.random as npr
import numpy.linalg as linalg

np.set_printoptions(precision=4, suppress=True)

print("Section6: Recurrence vs Transience (Monte Carlo heuristic)")
print("- Finite irreducible chain is positive recurrent: expected return time finite.")
# P3: 具有吸收状态的链（状态0为吸收态）
# 从状态1出发，在遇到吸收态之前返回的概率小于1
# 这表明状态1具有瞬态特性
P3 = np.array([
    [1.0, 0.0, 0.0],  # 状态0: 吸收态（100%停留在自己）
    [0.5, 0.5, 0.0],  # 状态1: 50%被吸收，50%停留在自己
    [0.0, 0.4, 0.6]  # 状态2: 40%转到状态1，60%停留在自己
])
# P5: 不可约的正常返链
# 估计的平均返回时间约为3
# 理论平稳概率约为0.333
# 验证了 π_i × m_i ≈ 1 的关
P5 = np.array([
    [0.5, 0.3, 0.2],
    [0.2, 0.5, 0.3],
    [0.3, 0.2, 0.5]
])


def estimate_return_time(P, start, trials=2000, max_steps=10000):
    times = []
    for _ in range(trials):
        cur = start
        for t in range(1, max_steps + 1):
            cur = npr.choice(a=len(P), p=P[cur])
            if cur == start:
                times.append(t)
                break
    return np.mean(times), np.median(times)


mean_rt, med_rt = estimate_return_time(P5, 0, 20000, 10000)
# 常返状态: 从该状态出发，几乎必然会在有限时间内返回
print(f"Estimated mean return time to state 0(P5):mean={mean_rt:.2f}, median={med_rt:.2f}")
print("Compare without theoretical: stationary s0=1/mean_return_time. Check s0*mean≈1")

A = P5.T - np.eye(3)  # A = (p^T - I) = -(I-P)^T
A2 = np.vstack([A, np.ones(3)])  # M=[A; 1^T]
b2 = np.array([0, 0, 0, 1.0])  # c=[0,0,0,1]^T
sol = linalg.lstsq(A2, b2, rcond=None)[0]
print("s0 from solve:", sol[0], "s0 * mean ≈", sol[0] * mean_rt)


# for a transient-like scenario(chain with absorbing states),
# starting in a non-absorbing state may never return
def estimate_return_probability(P, start, trials=2000, max_steps=10000):
    count = 0
    for _ in range(trials):
        cur = start
        returned = False
        for t in range(1, max_steps + 1):
            cur = npr.choice(a=len(P), p=P[cur])
            if cur == start:
                returned = True
                break
            if cur in {0}:  # absorbing 0 in P3
                break
        if returned:
            count += 1
    return count / trials


ret_prob = estimate_return_probability(P3, 1, trials=2000, max_steps=1000)
# 瞬态状态: 从该状态出发，有一定概率永远不会返回
print(f"Estimated probability of returning to state 1 (starting at 1) before absorption:"
      f"{ret_prob:.4f}(less than 1 -> transisent behaviour possible)")
